import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

a = np.loadtxt('200K_25meV.txt')
b = np.loadtxt('600K_25meV.txt')

factor1 = 0.4332306897319976 # normalized to calculated DOS
factor2 = 0.439281161781896 
x_mask = np.where((a[:,0] < 9) & (a[:,0] > 1.5)) 
plt.errorbar(a[x_mask,0][0], a[x_mask,1][0]*factor1, yerr=a[x_mask,2][0]*factor1,color='blue',marker='o',markersize=6,capsize=4,lw=2,label='200K INS')
plt.errorbar(b[x_mask,0][0], b[x_mask,1][0]*factor2, yerr=b[x_mask,2][0]*factor2,color='red',marker='o',markersize=6,capsize=4,lw=2,label='600K INS')
x0 = np.arange(0,6,0.1)
y0 = x0 * 0.0037122 + 0.00145116
y00 = x0 * 0.00141308 + 0.00016797
x1 = np.arange(0,9,0.1)
plt.plot(x0,y0*factor1,color='red',ls='dashed',lw=3)
plt.plot(x0,y00*factor2,color='blue',ls=(0, (1, 2)),lw=2)
plt.plot(x1,x1**2*4.91178158e-04*factor1,color='blue',ls='dashed',lw=3)
plt.plot(x1,x1**2*7.90805999e-04*factor2,color='red',ls=(0, (1, 2)),lw=2)
plt.plot(0,0.001034906914893617,label='_nolegend_',color='purple',marker='P',clip_on=False,markersize=15)
plt.text(0.1,0.004,'600K QENS',fontsize=15,color='purple')
plt.xlim(0,6)
plt.ylim(0,0.015)
plt.xticks([0,2,4,6])
plt.yticks([0.00,0.005,0.010,0.015])
plt.tick_params(which='both',direction='in',labelsize=20)
plt.minorticks_on()
plt.xlabel('Energy (meV)',fontsize=20)
plt.ylabel(r'DOS (meV$^{-1}$)',fontsize=20)
plt.legend(frameon=False,fontsize=20,loc='upper left',handlelength=1.5)
#plt.savefig('LowE.pdf',format='pdf',bbox_inches='tight',dpi=600)
#plt.savefig('LowE.png',format='png',bbox_inches='tight',dpi=600)
#plt.show()

linear = np.column_stack((x0, y0*factor1, y00*factor2))
quadratic = np.column_stack((x1,x1**2*4.91178158e-04*factor1,x1**2*7.90805999e-04*factor2))
np.savetxt('linear.txt',linear)
np.savetxt('quadratic.txt',quadratic)
